#!/usr/bin/env python3

import pyaudio
import json
import time
import os
import requests
import subprocess
import re
import threading
import tempfile
import signal
import sys
import logging
from datetime import datetime, timedelta
from vosk import Model, KaldiRecognizer

class RoboticArmVoiceControlDaemon:
    def __init__(self):
        # 设置日志
        self.setup_logging()
        
        # 初始化模型和API
        self.model = Model("../vosk-model-small-cn-0.22")
        self.api_key = "ms-fecb8eca-0503-44e9-9cf2-4870bb0ed37c"
        self.base_url = "https://api-inference.modelscope.cn/v1"
        self.model_name = "deepseek-ai/DeepSeek-V3.2-Exp"
        
        # 状态标志
        self.speaking = False
        self.recording = False
        self.running = True
        self.active_mode = False  # 是否处于激活状态
        self.last_activity_time = datetime.now()
        
        # 休眠设置 (30分钟无活动进入休眠)
        self.sleep_timeout = 1800  # 30分钟
        self.wake_word = "小臂"  # 唤醒词
        
        # 机械臂状态
        self.arm_state = {
            "base": 500,      # 1号舵机: 底座 (0-1000)
            "shoulder": 500,  # 2号舵机: 肩部 (0-1000)
            "elbow": 500,     # 3号舵机: 肘部 (0-1000)
            "wrist_vert": 500, # 4号舵机: 腕部垂直 (0-1000)
            "wrist_rot": 500, # 5号舵机: 腕部旋转 (0-1000)
            "gripper": 0      # 6号舵机: 夹爪 (0:闭合, 1000:打开)
        }
        
        # 舵机ID映射
        self.servo_ids = {
            "base": 1,
            "shoulder": 2,
            "elbow": 3,
            "wrist_vert": 4,
            "wrist_rot": 5,
            "gripper": 6
        }
        
        # 设置信号处理
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        
        self.log("🤖 机械臂语音控制守护程序已启动")
        self.log("📍 当前舵机状态: {}".format(self.arm_state))
        self.speak("系统已启动，等待唤醒")
        
    def setup_logging(self):
        """设置日志系统"""
        log_dir = "/var/log/robotic_arm"
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
            
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(os.path.join(log_dir, "robotic_arm.log")),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
        
    def log(self, message):
        """记录日志"""
        self.logger.info(message)
        print(message)
        
    def signal_handler(self, signum, frame):
        """处理退出信号"""
        self.log("收到退出信号，正在关闭...")
        self.running = False
        self.speak("系统关闭")
        sys.exit(0)
        
    def speak(self, text):
        """语音播报 - 使用espeak"""
        self.log("🔊 {}".format(text))
        
        if self.speaking:
            return
            
        def _speak():
            self.speaking = True
            try:
                with open(os.devnull, 'w') as devnull:
                    subprocess.run(["espeak", "-v", "zh", "-s", "200", text], 
                                 stdout=devnull, 
                                 stderr=devnull,
                                 timeout=5)
            except Exception as e:
                self.log("语音合成错误: {}".format(e))
            finally:
                self.speaking = False
        
        thread = threading.Thread(target=_speak)
        thread.daemon = True
        thread.start()
        
    def record_audio_direct(self, duration=5):
        """直接录制音频，不使用静默检测"""
        self.log("🎤 开始录音 {}秒...".format(duration))
        
        try:
            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
                temp_filename = temp_file.name
            
            with open(os.devnull, 'w') as devnull:
                result = subprocess.run([
                    "arecord", 
                    "-f", "S16_LE", 
                    "-r", "16000", 
                    "-c", "1", 
                    "-d", str(duration),
                    temp_filename
                ], stdout=devnull, stderr=devnull, timeout=duration+2)
            
            if result.returncode == 0:
                self.log("✅ 录音完成")
                return temp_filename
            else:
                self.log("❌ 录音失败")
                if os.path.exists(temp_filename):
                    os.unlink(temp_filename)
                return None
                
        except subprocess.TimeoutExpired:
            self.log("❌ 录音超时")
            if os.path.exists(temp_filename):
                os.unlink(temp_filename)
            return None
        except Exception as e:
            self.log("录音错误: {}".format(e))
            if os.path.exists(temp_filename):
                os.unlink(temp_filename)
            return None
    
    def transcribe_audio(self, audio_file):
        """转录音频文件"""
        if not audio_file or not os.path.exists(audio_file):
            return ""
            
        try:
            recognizer = KaldiRecognizer(self.model, 16000)
            
            with open(audio_file, 'rb') as f:
                audio_data = f.read()
            
            chunk_size = 4000
            text_result = ""
            
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i:i+chunk_size]
                if recognizer.AcceptWaveform(chunk):
                    result = recognizer.Result()
                    result_dict = json.loads(result)
                    if result_dict['text']:
                        text_result += result_dict['text'] + " "
            
            final_result = recognizer.FinalResult()
            final_dict = json.loads(final_result)
            if final_dict['text']:
                text_result += final_dict['text']
            
            os.unlink(audio_file)
            
            return text_result.strip()
            
        except Exception as e:
            self.log("转录错误: {}".format(e))
            if os.path.exists(audio_file):
                os.unlink(audio_file)
            return ""
    
    def check_wake_word(self, text):
        """检查是否包含唤醒词"""
        if self.wake_word in text:
            self.log("✅ 检测到唤醒词: {}".format(self.wake_word))
            return True
        return False
    
    def listen_for_wake_word(self):
        """监听唤醒词"""
        self.log("🔍 正在监听唤醒词: '{}'".format(self.wake_word))
        
        while self.running and not self.active_mode:
            # 检查是否应该进入休眠
            if self.should_sleep():
                self.enter_sleep_mode()
                continue
                
            # 录制短音频进行唤醒词检测
            audio_file = self.record_audio_direct(duration=3)
            if audio_file:
                text = self.transcribe_audio(audio_file)
                if text:
                    self.log("📝 监听识别: {}".format(text))
                    
                    if self.check_wake_word(text):
                        self.activate_system()
                        break
                    
                    # 更新活动时间
                    self.last_activity_time = datetime.now()
    
    def should_sleep(self):
        """检查是否应该进入休眠状态"""
        if self.active_mode:
            return False
            
        idle_time = (datetime.now() - self.last_activity_time).total_seconds()
        return idle_time >= self.sleep_timeout
    
    def enter_sleep_mode(self):
        """进入休眠模式"""
        self.log("💤 进入休眠模式...")
        self.speak("进入休眠模式，说{}唤醒".format(self.wake_word))
        
        # 在休眠模式下进行更高效的监听
        while self.running and not self.active_mode:
            audio_file = self.record_audio_direct(duration=2)
            if audio_file:
                text = self.transcribe_audio(audio_file)
                if text and self.check_wake_word(text):
                    self.activate_system()
                    break
    
    def activate_system(self):
        """激活系统"""
        self.active_mode = True
        self.last_activity_time = datetime.now()
        self.log("🎉 系统已激活")
        self.speak("我在，请下达指令")
        self.interactive_control()
    
    def analyze_voice_command(self, text):
        """分析语音命令并生成舵机控制指令"""
        system_prompt = """你是一个机械臂控制专家。请将用户的语音命令转换为具体的舵机控制指令。

机械臂有6个舵机：
1. base (底座): 0-1000位置，控制整体左右旋转
2. shoulder (肩部): 0-1000位置，控制大臂上下俯仰
3. elbow (肘部): 0-1000位置，控制小臂前后伸展
4. wrist_vert (腕部垂直): 0-1000位置，控制手腕上下摆动
5. wrist_rot (腕部旋转): 0-1000位置，控制手腕水平旋转
6. gripper (夹爪): 0-1000位置，0为闭合，1000为完全打开

请根据用户指令，生成JSON格式的控制指令，包含以下字段：
- action: 动作描述
- commands: 舵机控制列表，每个命令包含 servo(舵机名称) 和 position(目标位置，0-1000)
- duration: 执行时间(秒)

示例指令：
用户说："向右转"
输出：{"action": "底座向右旋转", "commands": [{"servo": "base", "position": 700}], "duration": 2}

用户说："抬起手臂"
输出：{"action": "抬起机械臂", "commands": [{"servo": "shoulder", "position": 300}, {"servo": "elbow", "position": 700}], "duration": 3}

用户说："抓取物体"
输出：{"action": "抓取物体", "commands": [{"servo": "gripper", "position": 0}], "duration": 1}

用户说："旋转手腕"
输出：{"action": "旋转手腕", "commands": [{"servo": "wrist_rot", "position": 800}], "duration": 2}

用户说："放下手腕"
输出：{"action": "放下手腕", "commands": [{"servo": "wrist_vert", "position": 300}], "duration": 2}

请只返回JSON格式，不要其他文字。"""

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"当前状态: {self.arm_state}。指令: {text}"}
            ],
            "temperature": 0.3,
            "max_tokens": 500
        }
        
        try:
            self.log("🔄 解析语音指令...")
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                response_text = result["choices"][0]["message"]["content"].strip()
                
                try:
                    command = json.loads(response_text)
                    return command
                except:
                    json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                    if json_match:
                        command = json.loads(json_match.group())
                        return command
                    else:
                        return {"error": "无法解析AI响应"}
            else:
                return {"error": f"API调用失败: {response.status_code}"}
                
        except Exception as e:
            return {"error": f"API调用异常: {e}"}
    
    def execute_commands(self, command_data):
        """执行舵机控制命令"""
        if "error" in command_data:
            error_msg = f"错误: {command_data['error']}"
            self.speak(error_msg)
            return False
        
        action = command_data.get('action', '未知动作')
        duration = command_data.get('duration', 2)
        
        self.speak(f"执行{action}")
        
        commands = command_data.get("commands", [])
        
        # 生成 movemulti 命令
        positions = []
        for servo_name in self.servo_ids.keys():
            target_pos = None
            for cmd in commands:
                if cmd.get("servo") == servo_name:
                    target_pos = cmd.get("position")
                    break
            
            if target_pos is not None:
                positions.append(str(target_pos))
                self.arm_state[servo_name] = target_pos
            else:
                positions.append("-")
        
        # 发送 movemulti 命令
        movemulti_cmd = f"movemulti {','.join(positions)}"
        self.log("[硬件控制] {}".format(movemulti_cmd))
        self.send_command_to_arm(movemulti_cmd)
        
        # 执行时间
        self.log("⏳ 执行中...", end="", flush=True)
        for i in range(duration):
            print(".", end="", flush=True)
            time.sleep(1)
        print(" 完成!")
        
        self.speak("完成")
        self.log("📍 更新后状态: {}".format(self.arm_state))
        
        # 更新活动时间
        self.last_activity_time = datetime.now()
        return True
    
    def control_single_servo(self, servo_name, position):
        """控制单个舵机"""
        if servo_name in self.servo_ids:
            servo_id = self.servo_ids[servo_name]
            command = f"move {servo_id} {position}"
            self.log("[硬件控制] {}".format(command))
            self.send_command_to_arm(command)
            self.arm_state[servo_name] = position
            
            # 更新活动时间
            self.last_activity_time = datetime.now()
    
    def send_command_to_arm(self, command):
        """发送命令到机械臂控制器"""
        self.log("[发送命令] {}".format(command))
        # 这里添加实际的硬件通信代码
    
    def voice_menu_selection(self, menu_options, prompt_text):
        """语音菜单选择"""
        while self.running:
            # 先显示菜单，再语音提示
            self.log("\n📋 可用选项:")
            for key, description in menu_options.items():
                self.log("  {}. {}".format(key, description))
            
            self.speak(prompt_text)
            time.sleep(1)
            
            # 录制音频
            audio_file = self.record_audio_direct(duration=5)
            if audio_file:
                # 转录音频
                choice = self.transcribe_audio(audio_file)
                self.log("📝 识别结果: {}".format(choice))
            else:
                choice = ""
            
            if choice:
                # 中文数字转阿拉伯数字
                chinese_to_arabic = {
                    '一': '1', '二': '2', '三': '3', '四': '4', '五': '5',
                    '六': '6', '七': '7', '八': '8', '九': '9', '十': '10',
                    '1': '1', '2': '2', '3': '3', '4': '4', '5': '5',
                    '6': '6', '7': '7', '8': '8', '9': '9'
                }
                
                # 直接匹配
                clean_choice = choice.replace(' ', '').strip()
                if clean_choice in chinese_to_arabic:
                    choice_num = chinese_to_arabic[clean_choice]
                else:
                    # 提取第一个数字
                    numbers = re.findall(r'\d', clean_choice)
                    if numbers:
                        choice_num = numbers[0]
                    else:
                        # 关键词匹配
                        for key, description in menu_options.items():
                            if key in clean_choice or description in clean_choice:
                                choice_num = key
                                break
                        else:
                            choice_num = None
                
                if choice_num and choice_num in menu_options:
                    selected = menu_options[choice_num]
                    self.log("✅ 选择: {}".format(selected))
                    self.speak(f"选择{selected}")
                    
                    # 更新活动时间
                    self.last_activity_time = datetime.now()
                    return choice_num
                
                # 检查取消词
                if any(word in choice for word in ["退出", "取消", "返回", "结束", "休眠"]):
                    return "cancel"
                
                self.log("❌ 无法识别: {}".format(choice))
                self.speak("请重新选择")
            else:
                self.speak("请说话")
    
    def interactive_control(self):
        """交互式控制模式"""
        self.speak("欢迎使用")
        
        while self.running and self.active_mode:
            # 检查是否应该返回休眠
            idle_time = (datetime.now() - self.last_activity_time).total_seconds()
            if idle_time > 300:  # 5分钟无活动返回监听模式
                self.speak("长时间无操作，返回监听模式")
                self.active_mode = False
                break
                
            menu_options = {
                "1": "语音控制",
                "2": "复位机械臂",
                "3": "读取状态",
                "4": "单个舵机控制",
                "5": "返回监听模式"
            }
            
            self.log("\n" + "="*40)
            self.log("🤖 机械臂语音控制系统")
            self.log("="*40)
            
            choice = self.voice_menu_selection(menu_options, "请选择")
            
            if choice == "1":
                self.voice_control_mode()
            elif choice == "2":
                self.reset_arm()
            elif choice == "3":
                self.read_servo_status()
            elif choice == "4":
                self.single_servo_control()
            elif choice == "5" or choice == "cancel":
                self.speak("返回监听模式")
                self.active_mode = False
                break
            else:
                self.speak("请重新选择")
    
    def voice_control_mode(self):
        """语音控制模式"""
        self.speak("请说话")
        audio_file = self.record_audio_direct(duration=5)
        
        if not audio_file:
            self.speak("没听到")
            return
        
        # 转录音频
        command = self.transcribe_audio(audio_file)
        
        if not command:
            self.speak("没听清")
            return
        
        self.log("🗣️  指令: {}".format(command))
        self.speak(f"收到{command}")
        
        # 解析指令
        command_data = self.analyze_voice_command(command)
        
        # 显示解析结果
        action = command_data.get('action', '未知动作')
        duration = command_data.get('duration', 2)
        
        self.log("🔍 动作: {}, 时长: {}秒".format(action, duration))
        self.speak(f"准备{action}")
        
        # 快速确认
        self.speak("确认吗")
        confirm_audio = self.record_audio_direct(duration=3)
        if confirm_audio:
            confirm_text = self.transcribe_audio(confirm_audio)
            confirm = confirm_text if confirm_text else ""
        else:
            confirm = ""
        
        if confirm and any(word in confirm for word in ["确认", "执行", "是的", "好", "开始"]):
            self.execute_commands(command_data)
        else:
            self.speak("取消")
    
    def reset_arm(self):
        """复位机械臂到初始位置"""
        self.speak("复位中")
        reset_commands = {
            "action": "复位机械臂",
            "commands": [
                {"servo": "base", "position": 500},
                {"servo": "shoulder", "position": 500},
                {"servo": "elbow", "position": 500},
                {"servo": "wrist_vert", "position": 500},
                {"servo": "wrist_rot", "position": 500},
                {"servo": "gripper", "position": 0}
            ],
            "duration": 2
        }
        self.execute_commands(reset_commands)
    
    def read_servo_status(self):
        """读取舵机状态"""
        status_text = "状态: "
        for servo_name, position in self.arm_state.items():
            servo_id = self.servo_ids.get(servo_name, "未知")
            status_text += f"{servo_id}号{position}, "
        
        self.log("📊 {}".format(status_text))
        self.speak(status_text)
    
    def single_servo_control(self):
        """单个舵机控制模式"""
        menu_options = {
            "1": "底座",
            "2": "肩部",
            "3": "肘部",
            "4": "腕部垂直", 
            "5": "腕部旋转",
            "6": "夹爪",
            "7": "返回"
        }
        
        choice = self.voice_menu_selection(menu_options, "选择舵机")
        
        servo_map = {
            "1": "base",
            "2": "shoulder",
            "3": "elbow", 
            "4": "wrist_vert",
            "5": "wrist_rot",
            "6": "gripper"
        }
        
        if choice in servo_map:
            servo_name = servo_map[choice]
            current_pos = self.arm_state[servo_name]
            
            self.speak(f"当前位置{current_pos}")
            pos_audio = self.record_audio_direct(duration=4)
            if pos_audio:
                pos_command = self.transcribe_audio(pos_audio)
            else:
                pos_command = ""
            
            if pos_command:
                numbers = re.findall(r'\d+', pos_command)
                if numbers:
                    pos = int(numbers[0])
                    if 0 <= pos <= 1000:
                        self.speak(f"设置{pos}")
                        self.control_single_servo(servo_name, pos)
                    else:
                        self.speak("范围错误")
                else:
                    self.speak("没听到数字")
        elif choice != "7" and choice != "cancel":
            self.speak("无效")
    
    def run_daemon(self):
        """运行守护程序主循环"""
        self.log("🚀 守护程序开始运行")
        
        try:
            while self.running:
                if not self.active_mode:
                    self.listen_for_wake_word()
                else:
                    time.sleep(0.1)  # 减少CPU占用
                    
        except KeyboardInterrupt:
            self.log("用户中断程序")
        except Exception as e:
            self.log("守护程序错误: {}".format(e))
        finally:
            self.log("守护程序停止")
            self.speak("系统关闭")

def main():
    # 检查是否以守护进程模式运行
    if len(sys.argv) > 1 and sys.argv[1] == "--daemon":
        # 守护进程模式
        daemon = RoboticArmVoiceControlDaemon()
        daemon.run_daemon()
    else:
        # 交互式模式
        controller = RoboticArmVoiceControlDaemon()
        controller.interactive_control()

if __name__ == "__main__":
    main()
